// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/container/f14_map.h>

#include <pollux/common/base/exceptions.h>
#include <pollux/common/memory/hash_string_allocator.h>
#include <pollux/exec/addressable_non_null_value_list.h>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/strings.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/functions/prestosql/aggregates/histogram_aggregate.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        template<
            typename T,
            typename Hash = std::hash<T>,
            typename EqualTo = std::equal_to<T> >
        struct Accumulator {
            using ValueMap = melon::F14FastMap<
                T,
                int64_t,
                Hash,
                EqualTo,
                AlignedStlAllocator<std::pair<const T, int64_t>, 16> >;

            ValueMap values;

            Accumulator(const TypePtr & /*type*/, HashStringAllocator *allocator)
                : values{
                    AlignedStlAllocator<std::pair<const T, int64_t>, 16>(allocator)
                } {
            }

            Accumulator(Hash hash, EqualTo equalTo, HashStringAllocator *allocator)
                : values{
                    0,
                    hash,
                    equalTo,
                    AlignedStlAllocator<std::pair<const T, int64_t>, 16>(allocator)
                } {
            }

            size_t size() const {
                return values.size();
            }

            void addValue(
                DecodedVector &decoded,
                vector_size_t index,
                HashStringAllocator * /*allocator*/) {
                values[decoded.value_at<T>(index)]++;
            }

            void addValueWithCount(
                T value,
                int64_t count,
                HashStringAllocator * /*allocator*/) {
                values[value] += count;
            }

            void extractValues(
                FlatVector<T> &keys,
                FlatVector<int64_t> &counts,
                vector_size_t offset) {
                auto index = offset;
                for (const auto &[value, count]: values) {
                    keys.set(index, value);
                    counts.set(index, count);

                    ++index;
                }
            }

            void free(HashStringAllocator &allocator) {
                using TValues = decltype(values);
                values.~TValues();
            }
        };

        // A wrapper around Accumulator that overrides hash and equal_to functions to
        // use the custom comparisons provided by a custom type.
        template<TypeKind Kind>
        struct CustomComparisonAccumulator {
            using NativeType = typename TypeTraits<Kind>::NativeType;

            struct Hash {
                const TypePtr &type;

                size_t operator()(const NativeType &value) const {
                    return static_cast<const CanProvideCustomComparisonType<Kind> *>(
                                type.get())
                            ->hash(value);
                }
            };

            struct EqualTo {
                const TypePtr &type;

                bool operator()(const NativeType &left, const NativeType &right) const {
                    return static_cast<const CanProvideCustomComparisonType<Kind> *>(
                               type.get())
                           ->compare(left, right) == 0;
                }
            };

            // The underlying Accumulator to which all operations are delegated.
            Accumulator<
                NativeType,
                CustomComparisonAccumulator::Hash,
                CustomComparisonAccumulator::EqualTo>
            base;

            CustomComparisonAccumulator(
                const TypePtr &type,
                HashStringAllocator *allocator)
                : base{
                    CustomComparisonAccumulator::Hash{type},
                    CustomComparisonAccumulator::EqualTo{type},
                    allocator
                } {
            }

            size_t size() const {
                return base.size();
            }

            void addValue(
                DecodedVector &decoded,
                vector_size_t index,
                HashStringAllocator *allocator) {
                base.addValue(decoded, index, allocator);
            }

            void addValueWithCount(
                NativeType value,
                int64_t count,
                HashStringAllocator *allocator) {
                base.addValueWithCount(value, count, allocator);
            }

            void extractValues(
                FlatVector<NativeType> &keys,
                FlatVector<int64_t> &counts,
                vector_size_t offset) {
                base.extractValues(keys, counts, offset);
            }

            void free(HashStringAllocator &allocator) {
                base.free(allocator);
            }
        };

        struct StringViewAccumulator {
            /// A map of unique StringViews pointing to storage managed by 'strings'.
            Accumulator<StringView> base;

            /// Stores unique non-null non-inline strings.
            Strings strings;

            StringViewAccumulator(const TypePtr &type, HashStringAllocator *allocator)
                : base{type, allocator} {
            }

            size_t size() const {
                return base.size();
            }

            void addValue(
                DecodedVector &decoded,
                vector_size_t index,
                HashStringAllocator *allocator) {
                auto value = decoded.value_at<StringView>(index);
                base.values[store(value, allocator)]++;
            }

            void addValueWithCount(
                StringView value,
                int64_t count,
                HashStringAllocator *allocator) {
                base.values[store(value, allocator)] += count;
            }

            StringView store(StringView value, HashStringAllocator *allocator) {
                if (!value.isInline()) {
                    auto it = base.values.find(value);
                    if (it != base.values.end()) {
                        value = it->first;
                    } else {
                        value = strings.append(value, *allocator);
                    }
                }
                return value;
            }

            void extractValues(
                FlatVector<StringView> &keys,
                FlatVector<int64_t> &counts,
                vector_size_t offset) {
                base.extractValues(keys, counts, offset);
            }

            void free(HashStringAllocator &allocator) {
                strings.free(allocator);
                using Base = decltype(base);
                base.~Base();
            }
        };

        struct ComplexTypeAccumulator {
            Accumulator<
                AddressableNonNullValueList::Entry,
                AddressableNonNullValueList::Hash,
                AddressableNonNullValueList::EqualTo>
            base;

            AddressableNonNullValueList serializedValues;

            ComplexTypeAccumulator(const TypePtr &type, HashStringAllocator *allocator)
                : base{
                    AddressableNonNullValueList::Hash{},
                    AddressableNonNullValueList::EqualTo{type},
                    allocator
                } {
            }

            size_t size() const {
                return base.size();
            }

            void addValue(
                DecodedVector &decoded,
                vector_size_t index,
                HashStringAllocator *allocator) {
                addValueWithCount(*decoded.base(), decoded.index(index), 1, allocator);
            }

            void addValueWithCount(
                const BaseVector &keys,
                vector_size_t index,
                int64_t count,
                HashStringAllocator *allocator) {
                auto entry = serializedValues.append(keys, index, allocator);

                auto it = base.values.find(entry);
                if (it == base.values.end()) {
                    // New entry.
                    base.values[entry] += count;
                } else {
                    // Existing entry.
                    serializedValues.removeLast(entry);

                    it->second += count;
                }
            }

            void extractValues(
                BaseVector &keys,
                FlatVector<int64_t> &counts,
                vector_size_t offset) {
                auto index = offset;
                for (const auto &[position, count]: base.values) {
                    AddressableNonNullValueList::read(position, keys, index);
                    counts.set(index, count);
                    ++index;
                }
            }

            void free(HashStringAllocator &allocator) {
                serializedValues.free(allocator);
                using Base = decltype(base);
                base.~Base();
            }
        };

        template<typename T>
        struct AccumulatorTypeTraits {
            using AccumulatorType = Accumulator<T>;
        };

        template<>
        struct AccumulatorTypeTraits<float> {
            using AccumulatorType = Accumulator<
                float,
                util::floating_point::NaNAwareHash<float>,
                util::floating_point::NaNAwareEquals<float> >;
        };

        template<>
        struct AccumulatorTypeTraits<double> {
            using AccumulatorType = Accumulator<
                double,
                util::floating_point::NaNAwareHash<double>,
                util::floating_point::NaNAwareEquals<double> >;
        };

        template<>
        struct AccumulatorTypeTraits<StringView> {
            using AccumulatorType = StringViewAccumulator;
        };

        template<>
        struct AccumulatorTypeTraits<ComplexType> {
            using AccumulatorType = ComplexTypeAccumulator;
        };

        // Combines a partial aggregation represented by the key-value pair at row in
        // mapKeys and mapValues into groupMap.
        template<typename T, typename Accumulator>
        MELON_ALWAYS_INLINE void addToFinalAggregation(
            const BaseVector *mapKeys,
            const FlatVector<int64_t> *flatValues,
            vector_size_t index,
            const vector_size_t *rawSizes,
            const vector_size_t *rawOffsets,
            Accumulator *accumulator,
            HashStringAllocator *allocator) {
            auto size = rawSizes[index];
            auto offset = rawOffsets[index];

            if constexpr (std::is_same_v<T, ComplexType>) {
                for (int i = 0; i < size; ++i) {
                    const auto count = flatValues->value_at(offset + i);
                    accumulator->addValueWithCount(*mapKeys, offset + i, count, allocator);
                }
            } else {
                auto *flatKeys = mapKeys->as_flat_vector<T>();
                for (int i = 0; i < size; ++i) {
                    accumulator->addValueWithCount(
                        flatKeys->value_at(offset + i),
                        flatValues->value_at(offset + i),
                        allocator);
                }
            }
        }

        template<
            typename T,
            typename AccumulatorType =
            typename AccumulatorTypeTraits<T>::AccumulatorType>
        class HistogramAggregate : public exec::Aggregate {
        public:
            explicit HistogramAggregate(TypePtr resultType)
                : Aggregate(std::move(resultType)) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(AccumulatorType);
            }

            bool isFixedSize() const override {
                return false;
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto *map_vector = (*result)->as<MapVector>();
                POLLUX_CHECK(map_vector);
                map_vector->resize(numGroups);

                auto &mapKeys = map_vector->mapKeys();
                auto *flatValues =
                        map_vector->mapValues()->as_unchecked<FlatVector<int64_t> >();
                POLLUX_CHECK_NOT_NULL(mapKeys);
                POLLUX_CHECK_NOT_NULL(flatValues);

                const auto numElements = countElements(groups, numGroups);
                mapKeys->resize(numElements);
                flatValues->resize(numElements);

                auto *rawNulls = map_vector->mutable_raw_nulls();
                vector_size_t offset = 0;
                if constexpr (std::is_same_v<T, ComplexType>) {
                    for (auto i = 0; i < numGroups; ++i) {
                        char *group = groups[i];
                        auto *accumulator = value<AccumulatorType>(group);

                        const auto mapSize = accumulator->size();
                        map_vector->setOffsetAndSize(i, offset, mapSize);
                        if (mapSize == 0) {
                            bits::set_null(rawNulls, i, true);
                        } else {
                            clearNull(rawNulls, i);
                            accumulator->extractValues(*mapKeys, *flatValues, offset);
                            offset += mapSize;
                        }
                    }
                } else {
                    auto *flatKeys = mapKeys->as_flat_vector<T>();
                    for (auto i = 0; i < numGroups; ++i) {
                        char *group = groups[i];
                        auto *accumulator = value<AccumulatorType>(group);

                        const auto mapSize = accumulator->size();
                        map_vector->setOffsetAndSize(i, offset, mapSize);
                        if (mapSize == 0) {
                            bits::set_null(rawNulls, i, true);
                        } else {
                            clearNull(rawNulls, i);
                            accumulator->extractValues(*flatKeys, *flatValues, offset);
                            offset += mapSize;
                        }
                    }
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                extractValues(groups, numGroups, result);
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                decodedKeys_.decode(*args[0], rows);

                rows.applyToSelected([&](auto row) {
                    if (decodedKeys_.is_null_at(row)) {
                        // Nulls among the values being aggregated are ignored.
                        return;
                    }
                    auto group = groups[row];
                    auto *accumulator = value<AccumulatorType>(group);

                    auto tracker = trackRowSize(group);
                    accumulator->addValue(decodedKeys_, row, allocator_);
                });
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /* mayPushdown */) override {
                decodedKeys_.decode(*args[0], rows);
                auto *accumulator = value<AccumulatorType>(group);

                auto tracker = trackRowSize(group);
                rows.applyToSelected([&](auto row) {
                    // Nulls among the values being aggregated are ignored.
                    if (!decodedKeys_.is_null_at(row)) {
                        accumulator->addValue(decodedKeys_, row, allocator_);
                    }
                });
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                decodedIntermediate_.decode(*args[0], rows);
                auto *indices = decodedIntermediate_.indices();
                auto *map_vector = decodedIntermediate_.base()->template as<MapVector>();

                auto &mapKeys = map_vector->mapKeys();
                auto *flatValues =
                        map_vector->mapValues()->template as_unchecked<FlatVector<int64_t> >();
                POLLUX_CHECK_NOT_NULL(mapKeys);
                POLLUX_CHECK_NOT_NULL(flatValues);

                auto rawSizes = map_vector->rawSizes();
                auto rawOffsets = map_vector->rawOffsets();
                rows.applyToSelected([&](vector_size_t row) {
                    if (!decodedIntermediate_.is_null_at(row)) {
                        auto group = groups[row];
                        auto *accumulator = value<AccumulatorType>(group);

                        auto tracker = trackRowSize(group);
                        addToFinalAggregation<T, AccumulatorType>(
                            mapKeys.get(),
                            flatValues,
                            indices[row],
                            rawSizes,
                            rawOffsets,
                            accumulator,
                            allocator_);
                    }
                });
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /* mayPushdown */) override {
                decodedIntermediate_.decode(*args[0], rows);
                auto *indices = decodedIntermediate_.indices();
                auto *map_vector = decodedIntermediate_.base()->template as<MapVector>();

                auto &mapKeys = map_vector->mapKeys();
                auto *flatValues =
                        map_vector->mapValues()->template as_unchecked<FlatVector<int64_t> >();
                POLLUX_CHECK_NOT_NULL(mapKeys);
                POLLUX_CHECK_NOT_NULL(flatValues);

                auto *accumulator = value<AccumulatorType>(group);

                auto tracker = trackRowSize(group);

                auto *rawSizes = map_vector->rawSizes();
                auto *rawOffsets = map_vector->rawOffsets();
                rows.applyToSelected([&](vector_size_t row) {
                    if (!decodedIntermediate_.is_null_at(row)) {
                        addToFinalAggregation<T, AccumulatorType>(
                            mapKeys.get(),
                            flatValues,
                            indices[row],
                            rawSizes,
                            rawOffsets,
                            accumulator,
                            allocator_);
                    }
                });
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                const auto &type = resultType()->childAt(0);
                for (auto index: indices) {
                    new(groups[index] + offset_) AccumulatorType{type, allocator_};
                }
            }

            void destroyInternal(melon::Range<char **> groups) override {
                for (auto *group: groups) {
                    if (isInitialized(group) && !isNull(group)) {
                        value<AccumulatorType>(group)->free(*allocator_);
                    }
                }
            }

        private:
            vector_size_t countElements(char **groups, int32_t numGroups) const {
                vector_size_t size = 0;
                for (int32_t i = 0; i < numGroups; ++i) {
                    size += value<AccumulatorType>(groups[i])->size();
                }
                return size;
            }

            DecodedVector decodedKeys_;
            DecodedVector decodedIntermediate_;
        };

        template<TypeKind Kind>
        std::unique_ptr<exec::Aggregate> createHistogramAggregateWithCustomCompare(
            const TypePtr &resultType) {
            return std::make_unique<HistogramAggregate<
                typename TypeTraits<Kind>::NativeType,
                CustomComparisonAccumulator<Kind> > >(resultType);
        }
    } // namespace

    void registerHistogramAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures = {
            exec::AggregateFunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("map(T,bigint)")
            .intermediateType("map(T,bigint)")
            .argumentType("T")
            .build(),
        };

        auto name = prefix + kHistogram;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(
                    argTypes.size(), 1, "{}: unexpected number of arguments", name);

                auto inputType =
                        exec::isRawInput(step) ? argTypes[0] : argTypes[0]->childAt(0);

                if (inputType->providesCustomComparison()) {
                    return POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                                createHistogramAggregateWithCustomCompare,
                                inputType->kind(),
                                resultType);
                }

                switch (inputType->kind()) {
                    case TypeKind::BOOLEAN:
                        return std::make_unique<HistogramAggregate<bool> >(resultType);
                    case TypeKind::TINYINT:
                        return std::make_unique<HistogramAggregate<int8_t> >(resultType);
                    case TypeKind::SMALLINT:
                        return std::make_unique<HistogramAggregate<int16_t> >(resultType);
                    case TypeKind::INTEGER:
                        return std::make_unique<HistogramAggregate<int32_t> >(resultType);
                    case TypeKind::BIGINT:
                        return std::make_unique<HistogramAggregate<int64_t> >(resultType);
                    case TypeKind::REAL:
                        return std::make_unique<HistogramAggregate<float> >(resultType);
                    case TypeKind::DOUBLE:
                        return std::make_unique<HistogramAggregate<double> >(resultType);
                    case TypeKind::TIMESTAMP:
                        return std::make_unique<HistogramAggregate<Timestamp> >(resultType);
                    case TypeKind::VARCHAR:
                    case TypeKind::VARBINARY:
                        return std::make_unique<HistogramAggregate<StringView> >(resultType);
                    case TypeKind::ARRAY:
                    case TypeKind::MAP:
                    case TypeKind::ROW:
                        return std::make_unique<HistogramAggregate<ComplexType> >(
                            resultType);
                    case TypeKind::HUGEINT:
                        if (inputType->isLongDecimal()) {
                            return std::make_unique<HistogramAggregate<int128_t> >(resultType);
                        }
                        [[fallthrough]];
                    case TypeKind::UNKNOWN:
                        return std::make_unique<HistogramAggregate<int8_t> >(resultType);
                    default:
                        POLLUX_NYI(
                            "Unknown input type for {} aggregation {}",
                            name,
                            inputType->toString());
                }
            },
            {false /*orderSensitive*/, false /*companionFunction*/},
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
